package limiter

import (
	"fmt"
	"sync"
	"time"
)

type Limiter struct {
	tb *TokenBucket
}

// 初始化限流器
func NewLimiter(r time.Duration, size int) *Limiter {
	return &Limiter{
		tb: &TokenBucket{
			rateLimit: r,
			size:      size,
			count:     size,
		},
	}
}

type TokenBucket struct {
	//锁
	sync.Mutex
	//桶大小
	size int
	//当前token数
	count int
	//填充速率,间隔多久补充token
	rateLimit time.Duration
	//最后请求成功的时间
	LastReqTime time.Time
}

func (tb *TokenBucket) fillToken() {
	tb.count += tb.getFillTokenCount()
}

func (tb *TokenBucket) getFillTokenCount() int {
	if tb.count >= tb.size {
		return 0
	}

	//如果不是第一次请求
	if !tb.LastReqTime.IsZero() {
		duration := time.Now().Sub(tb.LastReqTime)
		count := int(duration / tb.rateLimit)
		fmt.Println("duration=", duration)
		fmt.Println("rateLimit=", tb.rateLimit)
		fmt.Println("count=", count)
		//最多填满桶大小size
		if tb.size-tb.count >= count {
			return count
		} else {
			return tb.size - tb.count
		}
	}
	return 0
}

func (tb *TokenBucket) Allow() bool {
	//填充token数
	tb.fillToken()
	//减少token数
	if tb.count > 0 {
		tb.count--
		tb.LastReqTime = time.Now()
		return true
	}

	return false
}

func (l *Limiter) Allow() bool {
	l.tb.Lock()
	defer l.tb.Unlock()
	//计算补充token
	//当前token是否满足本次消耗
	return l.tb.Allow()
}
